from paddleocr import PaddleOCR, draw_ocr
import cv2
import numpy as np
from PIL import Image
import os

def unsharp_mask(image_path, output_path, kernel_size=(5, 5), sigma=1.0, amount=1.0, threshold=0):
    """非锐化掩模实现图像锐化
    
    参数:
        image_path: 输入图像路径
        output_path: 输出图像路径
        kernel_size: 高斯核大小
        sigma: 高斯核标准差
        amount: 锐化强度(0-2之间)
        threshold: 锐化阈值
    """
    # 读取图像
    image = cv2.imread(image_path)
    
    # 高斯模糊
    blurred = cv2.GaussianBlur(image, kernel_size, sigma)
    
    # 计算锐化掩模
    sharpened = float(amount + 1) * image - float(amount) * blurred
    
    # 确保像素值在0-255范围内
    sharpened = np.maximum(sharpened, np.zeros(sharpened.shape))
    sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape))
    sharpened = sharpened.round().astype(np.uint8)
    
    # 如果设置了阈值，只锐化高于阈值的区域
    if threshold > 0:
        low_contrast_mask = abs(image - blurred) < threshold
        np.copyto(sharpened, image, where=low_contrast_mask)
    
    # 保存结果
    cv2.imwrite(output_path, sharpened)
    
    return sharpened

def sort_ocr_results(ocr_result):
    """
    对OCR结果进行排序：先按x坐标由小到大，再按y坐标由小到大
    
    参数:
        ocr_result: OCR识别结果
    
    返回:
        排序后的OCR结果
    """
    # 计算每个文本框的中心点坐标
    sorted_results = []
    for idx in range(len(ocr_result)):
        res = ocr_result[idx]
        for line in res:
            # 获取文本框的四个顶点坐标
            box = line[0]
            # 计算中心点坐标
            center_x = sum(point[0] for point in box) / 4
            center_y = sum(point[1] for point in box) / 4
            # 添加到结果列表，包含中心点坐标和原始数据
            sorted_results.append({
                'center_x': center_x,
                'center_y': center_y,
                'data': line
            })
    
    # 按中心点坐标排序：先x后y
    sorted_results.sort(key=lambda item: (item['center_x'], item['center_y']))
    
    # 提取排序后的原始数据
    return [item['data'] for item in sorted_results]

# 初始化PaddleOCR
ocr = PaddleOCR(
    use_angle_cls=False, 
    lang='ch',
    rec_image_shape="3, 64, 512",
    det_limit_side_len=2500,
    det_limit_type="max",
    det_db_unclip_ratio=2.5,
    use_dilation=True,
    det_db_score_mode="slow",
    rec_thresh=0.7
)

# 图像锐化处理
sharpened_image = unsharp_mask(
    'D:/ocr/sample/1.jpg', 
    'D:/ocr/sharp.jpg',  
    kernel_size=(5, 5), 
    sigma=1.0, 
    amount=1.0, 
    threshold=0
)

# OCR识别
img_path = 'D:/ocr/sharp.jpg'
result = ocr.ocr(img_path, cls=False)

# 对结果进行排序
sorted_result = sort_ocr_results(result)

# 打印排序后的结果
print("排序后的OCR识别结果:")
for line in sorted_result:
    print(f"文本: {line[1][0]}, 置信度: {line[1][1]:.2f}, 中心坐标: ({sum(point[0] for point in line[0])/4:.1f}, {sum(point[1] for point in line[0])/4:.1f})")

# 绘制结果（使用排序后的结果）
image = Image.open(img_path).convert('RGB')
boxes = [line[0] for line in sorted_result]
txts = [line[1][0] for line in sorted_result]
scores = [line[1][1] for line in sorted_result]

# 绘制OCR结果
im_show = draw_ocr(image, boxes, txts, scores, font_path='/path/to/PaddleOCR/doc/fonts/simfang.ttf')
im_show = Image.fromarray(im_show)
im_show.save('result.jpg')
print("结果已保存到result.jpg")